import pickle

import pandas as pd

# 读取数据
excess_return = pd.read_excel(r'../data/Multifactor_Model_Data/1异象回归用超额收益率.xlsx')
momentum_data = pd.read_excel(r'../data/Multifactor_Model_Data/3动量效应.xlsx')
reverse_data = pd.read_excel(r'../data/Multifactor_Model_Data/3反转效应.xlsx')
size_data = pd.read_excel(r'../data/Multifactor_Model_Data/3总市值.xlsx')
BM_data = pd.read_excel(r'../data/Multifactor_Model_Data/3BM.xlsx')
close_price = pd.read_excel(r'../data/Multifactor_Model_Data/2全部股价序列（收盘价）.xlsx')
frames = [excess_return, momentum_data, reverse_data, size_data, BM_data, close_price]
for frame in frames:
    frame['time'] = frame['time'].astype(str)
    frame.set_index('time', drop=True, append=False, inplace=True)

data_dict = {}
time_list = excess_return.index.values.tolist()
for time in time_list:
    tmp_df = pd.DataFrame()
    tmp_df['excess_return'] = excess_return.loc[time]
    tmp_df['momentum'] = momentum_data.loc[time]
    tmp_df['reverse'] = reverse_data.loc[time]
    tmp_df['size'] = size_data.loc[time]
    tmp_df['BM'] = BM_data.loc[time]
    tmp_df['close'] = close_price.loc[time]
    data_dict[time] = tmp_df

with open("../data/test.pkl", "wb") as fp:  # Pickling
    pickle.dump(data_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)

# 清洗数据
with open("../data/test.pkl", 'rb') as file:
    raw_data = pickle.load(file)

date_list = sorted(raw_data.keys())
index_list = raw_data[date_list[0]].index.values

for date in date_list:
    cleaned_data_frame = raw_data[date].dropna(axis=0)
    index_list = list(set(index_list).intersection(set(cleaned_data_frame.index.values)))

cleaned_dict = {}
for key, value in raw_data.items():
    cleaned_data_frame = value[value.index.isin(index_list)]
    cleaned_dict[key] = cleaned_data_frame

with open("../data/weekly_data.pkl", "wb") as fp:  # Pickling
    pickle.dump(cleaned_dict, fp, protocol=pickle.HIGHEST_PROTOCOL)
